import torch
import seaborn
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# LSTM类
class LSTM(torch.nn.Module):
    def __init__(self, input_size=1, hidden_layer_size=100, output_size=1):
        super().__init__()
        self.hidden_layer_size = hidden_layer_size
        self.lstm = torch.nn.LSTM(input_size, hidden_layer_size)
        self.add_module('linear',torch.nn.Linear(100,1))

    def forward(self, input_seq):
        lstm_out, self.hidden_cell = self.lstm(input_seq.view(len(input_seq), 1, -1), self.hidden_cell)
        predictions = self.linear(lstm_out.view(len(input_seq), -1))
        return predictions[-1]


def SetBatch(Seq, L):  # 分批次函数
    Res = []
    Length = len(Seq)
    for i in range(Length - L):
        data = Seq[i:i + L]
        label = Seq[i + L]
        Res.append((data, label))
    return Res


# 获取数据集
FlightData = seaborn.load_dataset("flights")
# 得到乘客序列
Sequence = FlightData['passengers'].values.astype(float)
# 标准化
scaler = MinMaxScaler(feature_range=(-1, 1))
Sequence= scaler.fit_transform(Sequence.reshape(-1, 1))
# 转化为Tensor变量,便于接下来训练
Sequence = torch.FloatTensor(Sequence).view(-1)
# 切分数据集得到训练集与验证集
TrainData = Sequence[0:100]
ValidatingData = Sequence[100:]
# 设置模型
LSTM_Model = LSTM()
loss_function = torch.nn.MSELoss()
optimizer = torch.optim.Adam(LSTM_Model.parameters(), lr=0.001)
# 分批
ProcessedData = SetBatch(TrainData, 60)

# 开始训练
# 设置训练次数
epochs = 300

for i in range(epochs):
    # 每次训练过程中遍历ProcessedData
    for seq, labels in ProcessedData:
        optimizer.zero_grad()
        LSTM_Model.hidden_cell = (torch.zeros(1, 1, LSTM_Model.hidden_layer_size),
                                  torch.zeros(1, 1, LSTM_Model.hidden_layer_size))
        # 预测值
        y_prediction = LSTM_Model.forward(seq)
        # 计算损失
        single_loss = loss_function(y_prediction, labels)
        single_loss.backward()
        optimizer.step()
    # 观察Loss情况
    if i % 10 == 0:
        print(f'epoch: {i:3} loss: {single_loss.item():10.8f}')

LSTM_Model.eval()
# 训练完毕,进行测试
# 初始化预测序列
PredictedSeq = list(TrainData)

while len(PredictedSeq) != len(Sequence):
    seq = torch.FloatTensor(PredictedSeq[-60:])
    with torch.no_grad():
        LSTM_Model.hidden = (torch.zeros(1, 1, LSTM_Model.hidden_layer_size),
                        torch.zeros(1, 1, LSTM_Model.hidden_layer_size))
    PredictedVal = LSTM_Model.forward(seq).item()
    PredictedSeq.append(PredictedVal)

PredictedSeq = np.array(PredictedSeq)

plt.plot(PredictedSeq)
plt.plot(Sequence)

plt.show()